import redis
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Optional

class SessionManager:
    """会话管理器，用于支持多轮对话"""
    
    def __init__(self, redis_url: str, session_timeout: int = 3600):
        """
        初始化会话管理器
        
        Args:
            redis_url: Redis连接URL
            session_timeout: 会话超时时间（秒）
        """
        try:
            self.redis_client = redis.from_url(redis_url)
            self.redis_client.ping()  # 测试连接
        except Exception as e:
            print(f"Redis连接失败，使用内存存储: {e}")
            self.redis_client = None
            self._memory_store = {}
        
        self.session_timeout = session_timeout
    
    def create_session(self) -> str:
        """创建新会话"""
        session_id = str(uuid.uuid4())
        session_data = {
            'session_id': session_id,
            'created_at': datetime.now().isoformat(),
            'last_activity': datetime.now().isoformat(),
            'conversation_history': [],
            'context': {}
        }
        
        self._store_session(session_id, session_data)
        return session_id
    
    def get_session(self, session_id: str) -> Optional[Dict]:
        """获取会话数据"""
        if not session_id:
            return None
        
        session_data = self._get_session(session_id)
        if not session_data:
            return None
        
        # 检查会话是否过期
        last_activity = datetime.fromisoformat(session_data['last_activity'])
        if datetime.now() - last_activity > timedelta(seconds=self.session_timeout):
            self.delete_session(session_id)
            return None
        
        return session_data
    
    def update_session(self, session_id: str, data: Dict) -> bool:
        """更新会话数据"""
        session_data = self.get_session(session_id)
        if not session_data:
            return False
        
        session_data.update(data)
        session_data['last_activity'] = datetime.now().isoformat()
        
        self._store_session(session_id, session_data)
        return True
    
    def add_to_conversation(self, session_id: str, user_input: str, ai_response: str, 
                          input_image: Optional[str] = None, output_image: Optional[str] = None) -> bool:
        """添加对话记录到会话"""
        session_data = self.get_session(session_id)
        if not session_data:
            return False
        
        conversation_entry = {
            'timestamp': datetime.now().isoformat(),
            'user_input': user_input,
            'ai_response': ai_response,
            'input_image': input_image,
            'output_image': output_image
        }
        
        session_data['conversation_history'].append(conversation_entry)
        session_data['last_activity'] = datetime.now().isoformat()
        
        self._store_session(session_id, session_data)
        return True
    
    def get_conversation_history(self, session_id: str, limit: int = 10) -> List[Dict]:
        """获取对话历史"""
        session_data = self.get_session(session_id)
        if not session_data:
            return []
        
        history = session_data.get('conversation_history', [])
        return history[-limit:] if limit > 0 else history
    
    def delete_session(self, session_id: str) -> bool:
        """删除会话"""
        try:
            if self.redis_client:
                self.redis_client.delete(f"session:{session_id}")
            else:
                self._memory_store.pop(session_id, None)
            return True
        except Exception as e:
            print(f"删除会话失败: {e}")
            return False
    
    def _store_session(self, session_id: str, data: Dict):
        """存储会话数据"""
        try:
            if self.redis_client:
                self.redis_client.setex(
                    f"session:{session_id}",
                    self.session_timeout,
                    json.dumps(data, ensure_ascii=False)
                )
            else:
                self._memory_store[session_id] = data
        except Exception as e:
            print(f"存储会话失败: {e}")
    
    def _get_session(self, session_id: str) -> Optional[Dict]:
        """获取会话数据"""
        try:
            if self.redis_client:
                data = self.redis_client.get(f"session:{session_id}")
                return json.loads(data) if data else None
            else:
                return self._memory_store.get(session_id)
        except Exception as e:
            print(f"获取会话失败: {e}")
            return None
